from __future__ import with_statement
from khronos import *
from khronos.extras.components import *

shows = "Soap opera", "Movie", "News", "Cartoon", "Quiz show"
durations = HOUR, 45 * MINUTE, 30 * MINUTE

class TvStation(Atomic):
    default_priority = 1.0
    viewers = TSeriesAttr(storing=True, timescale=DAY)
    
    def __init__(self, color=None, *args, **kwargs):
        Atomic.__init__(self, *args, **kwargs)
        TvStation.viewers.initialize(self)
        self.color = color
        
    def status(self):
        return "%s viewers" % (self.viewers,)
        
    @chain
    def initializer(self):
        TvStation.viewers.reset(self, 0)
        rng = self.get_rng()
        while self.get_timedelta() < WEEK:
            # broadcast a random show
            with self.broadcasting(rng.choice(shows)):
                yield rng.choice(durations) + rng.gauss(0, 5) * MINUTE
            # giv'em commercials! oh yeah!
            with self.broadcasting("Commercials"):
                yield max(0, rng.gauss(10, 5)) * MINUTE
                
    def finalizer(self):
        viewers = TvStation.viewers.stat(self)
        try:
            comp = self.get_parent().viewers
        except AttributeError:
            comp = self.get_parent().viewers = Plotter(title="TV example", rows=3, cols=2)
        finally:
            viewers.run_chart(plotter=comp, title=self.get_name(), color=self.color)
            viewers.histogram(plotter=comp, title=self.get_name(), color=self.color)
            
class Viewers(Atomic):
    def __init__(self, initial=10, *args, **kwargs):
        Atomic.__init__(self, *args, **kwargs)
        self.initial = initial
        
    @chain
    def initializer(self):
        printer = self.launch(DatePrinter(period=12*HOUR))
        
        x = 0
        while x < self.initial:
            self.launch(Viewer("v%d" % (x,)))
            x += 1
            
        while self.get_timedelta() < WEEK:
            yield self.get_rng().gauss(10, 2) * MINUTE
            self.launch(Viewer("v%d" % (x,)))
            x += 1
        printer.terminate()
        
class Viewer(Atomic):
    @chain
    def initializer(self):
        rng = self.get_rng()
        preferred_show = rng.choice(shows)
        switch_channel_probability = 0.5
        
        stations = [self.get_root()[station] for station in ("RTP", "SIC", "TVI")]
        station = rng.choice(stations)
        station.viewers += 1
        while True:
            yield Listener([preferred_show + " start"])
            new_station = self.get_context(Signal).origin
            if self.__switch(station, new_station):
                station = new_station
                
            yield Listener([preferred_show + " finish"], targets=[station])
            if rng.random() < switch_channel_probability:
                while True:
                    new_station = rng.choice(stations)
                    if self.__switch(station, new_station):
                        station = new_station
                        break
                        
    def __switch(self, s1, s2):
        if s2 is s1:
            return False
        else:
            s1.viewers -= 1
            s2.viewers += 1
            return True
            
def build_model():
    root = Composite("root", children=[TvStation("blue",  "RTP"), 
                                       TvStation("green", "SIC"), 
                                       TvStation("red",   "TVI"), 
                                       Viewers(100, "viewers")])
    return root
    
def build():
    k = Khronos(build_model())
    return k, k.model
    
def main():
    k, m = build()
    k.single_run(trace=False)
    
k, m = build()